import os

from detectron2.data import DatasetCatalog, MetadataCatalog
from detectron2.data.datasets.builtin_meta import _get_builtin_metadata
from detectron2.data.datasets.cityscapes import load_cityscapes_semantic

from maskdino.data.datasets.dianjiao import load_dianjiao_semantic

# ==== Predefined splits for raw cityscapes images ===========
_RAW_CITYSCAPES_SPLITS = {
    "cityscapes_custom_{task}_train": ("cityscapes/leftImg8bit/train/", "cityscapes/gtFine/train/"),
    "cityscapes_custom_{task}_val": ("cityscapes/leftImg8bit/val/", "cityscapes/gtFine/val/"),
    "cityscapes_custom_{task}_test": ("cityscapes/leftImg8bit/test/", "cityscapes/gtFine/test/"),
}

_RAW_DIANJIAO_SPLITS = {
    "dianjiao_custom_{task}_train": ("workspaces/hjh9_data/dianjiao_dataset/Geneva_4G_0824", "workspaces/hjh9_data/dianjiao_dataset/Geneva_4G_0824"),
    # "dianjiao_custom_{task}_train": ("/dfs/data/workspaces/MaskDINO/tr_dataset_0718.csv", "/dfs/data/workspaces/MaskDINO/tr_dataset_0718.csv"),
    "dianjiao_custom_{task}_val": ("workspaces/dianjiao/", "workspaces/dianjiao/"),
    # "dianjiao_custom_{task}_val": ("/dfs/data/workspaces/MaskDINO/validate.csv", "/dfs/data/workspaces/MaskDINO/validate.csv"),
}

root = "workspaces/hjh9_data"

def register_all_cityscapes(root):
    for key, (image_dir, gt_dir) in _RAW_CITYSCAPES_SPLITS.items():
        meta = _get_builtin_metadata("cityscapes")
        image_dir = os.path.join(root, image_dir)
        gt_dir = os.path.join(root, gt_dir)

        sem_key = key.format(task="sem_seg")
        DatasetCatalog.register(
            sem_key, lambda x=image_dir, y=gt_dir: load_cityscapes_semantic(x, y)
        )
        MetadataCatalog.get(sem_key).set(
            image_dir=image_dir,
            gt_dir=gt_dir,
            evaluator_type="cityscapes_sem_seg",
            ignore_label=255,
            **meta,
        )

def register_all_dianjiao(root):
    for key, (image_dir, gt_dir) in _RAW_DIANJIAO_SPLITS.items():
        DIANJIAO_THING_CLASSES = [
            "jiaolu",
        ]
        DIANJIAO_STUFF_CLASSES = [
            "jiaolu",
        ]
        # fmt: on
        meta = {
            "thing_classes": DIANJIAO_THING_CLASSES,
            "stuff_classes": DIANJIAO_STUFF_CLASSES,
        }
        image_dir = os.path.join(root, image_dir)
        gt_dir = os.path.join(root, gt_dir)

        sem_key = key.format(task="sem_seg")
        DatasetCatalog.register(
            sem_key, lambda x=image_dir, y=gt_dir: load_dianjiao_semantic(x, y)
        )
        MetadataCatalog.get(sem_key).set(
            image_dir=image_dir,
            gt_dir=gt_dir,
            evaluator_type="dianjiao_sem_seg",
            ignore_label=255,
            **meta,
        )

# register_all_cityscapes(root)
register_all_dianjiao("")